#ifndef TCP_SERVER_H
#define TCP_SERVER_H

#include <ctime>
#include <functional>
#include <iostream>
#include <memory>
#include <string>
#include <asio.hpp>
#include <queue>
#include "io_context_pool.h"
#include "spin_lock.hpp"

class TcpConnection : public std::enable_shared_from_this<TcpConnection>
{
public:
    class SendBuff
    {
    public:
        explicit SendBuff(uint32_t buff_len) : buff_len_(buff_len)
        {
            buff_ = new char[buff_len_];
        }

        ~SendBuff()
        {
            delete[] buff_;
        }

        uint32_t GetBuffLen() const
        {
            return buff_len_;
        }

        char* GetBuff()
        {
            return buff_;
        }

        void SetDataLen(uint32_t data_len)
        {
            data_len_ = data_len;
        }

        uint32_t GetDataLen() const
        {
            return data_len_;
        }

        void SetSendLen(uint32_t send_len)
        {
            send_len_ = send_len;
        }

        uint32_t GetSendLen() const
        {
            return send_len_;
        }

    private:
        char* buff_ = nullptr;
        uint32_t buff_len_ = 0;
        uint32_t data_len_ = 0;
        uint32_t send_len_ = 0;
    };

    TcpConnection(asio::io_context &io_context, asio::ip::tcp::socket socket,
                  uint32_t recv_buff_len, uint32_t interval_time)
        : io_context_(io_context), socket_(std::move(socket)), recv_buff_len_(recv_buff_len),
          interval_time_(interval_time),
          send_timer_(io_context_, asio::chrono::milliseconds(interval_time_))
    {
        recv_buff_ = new char[recv_buff_len_];
    }

    virtual ~TcpConnection()
    {
        Close();
        delete[] recv_buff_;
    }

    void PushSendBuff(const std::shared_ptr<SendBuff> &send_buff)
    {
        send_queue_lock_.Lock();
        send_queue_.emplace(send_buff);
        send_queue_lock_.Unlock();
    }

    virtual void Connected() {}

    void Start()
    {
        DoRecv();
        DoSend();
    }

    virtual uint32_t ProcessRecv(char* data, uint32_t data_len) {}

    virtual void RecvError(const std::error_code &error, size_t bytes_transferred) {}

    virtual void SendError(const std::error_code &error, size_t bytes_transferred) {}

private:
    void DoRecv()
    {
        if (stop_)
        {
            return;
        }

        auto self(shared_from_this());
        socket_.async_read_some(
            asio::buffer(recv_buff_ + recv_data_len_, ((recv_buff_len_ - recv_data_len_) > 0)
                                                          ? (recv_buff_len_ - recv_data_len_)
                                                          : 0),
            [this, self](const std::error_code &error, size_t bytes_transferred) {
                if (!error)
                {
                    recv_data_len_ = recv_data_len_ + bytes_transferred;
                    const uint32_t process_len = ProcessRecv(recv_buff_, recv_data_len_);
                    recv_data_len_ = recv_data_len_ - process_len;
                    memmove(recv_buff_, recv_buff_ + process_len, recv_data_len_);
                }
                else
                {
                    RecvError(error, bytes_transferred);
                    stop_ = true;
                }
                DoRecv();
            });
    }

    void DoSend()
    {
        if (stop_)
        {
            return;
        }

        auto self(shared_from_this());
        std::shared_ptr<SendBuff> send_buff;
        if (GetSendBuff(send_buff))
        {
            asio::async_write(
                socket_,
                asio::buffer(send_buff->GetBuff() + send_buff->GetSendLen(),
                             send_buff->GetDataLen() - send_buff->GetSendLen()),
                [send_buff, this, self](const std::error_code &error, size_t bytes_transferred) {
                    if (!error)
                    {
                        send_buff->SetSendLen(send_buff->GetSendLen() + bytes_transferred);
                    }
                    else
                    {
                        SendError(error, bytes_transferred);
                        stop_ = true;
                    }
                    DoSend();
                });
        }
        else
        {
            send_timer_.async_wait([this, self](const std::error_code &error) {
                if (!error)
                {
                    send_timer_.expires_after(asio::chrono::milliseconds(interval_time_));
                    DoSend();
                }
                else
                {
                    /* code */
                }
            });
        }
    }

    bool GetSendBuff(std::shared_ptr<SendBuff> &send_buff)
    {
        send_queue_lock_.Lock();

        while (!send_queue_.empty())
        {
            const std::shared_ptr<SendBuff> &send_buff_ref = send_queue_.front();
            if (send_buff_ref->GetSendLen() >= send_buff_ref->GetDataLen())
            {
                send_queue_.pop();
            }
            else
            {
                send_buff = send_buff_ref;
                send_queue_lock_.Unlock();
                return true;
            }
        }

        send_queue_lock_.Unlock();
        return false;
    }

    void Close()
    {
        std::error_code ignored_ec;
        socket_.shutdown(asio::ip::tcp::socket::shutdown_both, ignored_ec);
        socket_.close(ignored_ec);
    }

    std::queue<std::shared_ptr<SendBuff>> send_queue_;
    SpinLock send_queue_lock_;
    asio::io_context &io_context_;
    uint32_t interval_time_;
    asio::steady_timer send_timer_;
    asio::ip::tcp::socket socket_;
    char* recv_buff_ = nullptr;
    uint32_t recv_buff_len_ = 0;
    uint32_t recv_data_len_ = 0;
    bool stop_ = false;
};

class TcpServer
{
public:
    TcpServer(IoContextPool &io_context_pool, const std::string &ip, int32_t port)
        : io_context_pool_(io_context_pool), acceptor_(io_context_pool_.GetIoContext())
    {
        const asio::ip::tcp::endpoint endpoint(asio::ip::address::from_string(ip), port);
        acceptor_.open(endpoint.protocol());
        acceptor_.set_option(asio::ip::tcp::acceptor::reuse_address(true));
        acceptor_.bind(endpoint);
        acceptor_.listen();
        DoAccept();
    }

private:
    void DoAccept()
    {
        acceptor_.async_accept(io_context_pool_.GetIoContext(),
                               [this](std::error_code ec, asio::ip::tcp::socket socket) {
                                   if (!ec)
                                   {
                                       std::make_shared<TcpConnection>(std::move(socket))->start();
                                   }
                                   else
                                   {
                                       std::cout << "async_accept error: " << ec.message()
                                                 << " code: " << ec.value() << '\n';
                                   }

                                   DoAccept();
                               });
    }

    IoContextPool &io_context_pool_;
    asio::ip::tcp::acceptor acceptor_; // acceptor_.close();
};

#endif